import os
import torch
import torch.nn.functional as F
import torch.nn as nn
from seg_utils import data_loader
import numpy as np
from PIL import Image
from tqdm import tqdm
from network import model
#from network.seg_hrnet import HighResolutionNet
#from network import deeplabv3
from network import model
#from network import unet
#from network import pspnet
batch_size =1
img_size = 512
def getLastCheckPt(model_path):
    files=os.listdir(model_path)
    max=0
    for file in files:
        if file[-4:]!=".pth":
            continue

        epoch,Extension =file.split(".")
        epoch=int(epoch)
        
        if max< epoch:
            max=epoch
    return max+1 #新的周期数
test_data = 'D:\\WHUData\\航空瓦片\\'
model_name = 'CFENet预训练'

save_path = 'D:\\WHUData\\航空瓦片\\'+model_name+"predictResult"
fileName=getLastCheckPt(test_data+model_name)
if fileName>1:
    model_path = 'D:\\WHUData\\航空瓦片\\'+model_name+"\\"+str(fileName-1)+".pth"
if not os.path.exists(save_path):
    os.makedirs(save_path)

def main():
    with torch.no_grad():
        # net = unet.unet_resnet101(n_classes=1, batch_size=batch_size, pretrained=True, fixed_feature=True).to(device)
        # net = deeplabv3.DeepLabV3().to(device)
        # net = HighResolutionNet(input_channels=3, output_channels=1).to(device)
        # net = pspnet.pspnet_resnet101(n_classes=1, batch_size =batch_size , pretrained=True, fixed_feature=True).to(device)
        net = model.CLNet().to(device)
        net.load_state_dict(torch.load(model_path, map_location='cuda:0'))
        perdict_loader = data_loader.get_loader(test_data+"predict\\" , batch_size , img_size,num_workers=4, mode='predict',augmentation_prob=0,shuffle=False, pin_memory=True)

        print("Strat predict!")
        predict(perdict_loader, net, save_path)
def predict(dataLoader, net, save_path):
    net.train(False)
    net.eval()
    for i, (inputs, filename) in enumerate(tqdm(dataLoader)):#每个批
        X = inputs.to(device)
        #Y = mask.to(device)
        output = net(X)
        # output = net(X)
        output = F.sigmoid(output)
        for i in range(output.shape[0]):#对每个批的每个图像识别结果分开处理
            probs_array = (torch.squeeze(output[i])).data.cpu().numpy()
            mask_array = (probs_array > 0.5)
            final_mask = mask_array.astype(np.float32)
            final_mask = final_mask * 255
            final_mask = final_mask.astype(np.uint8)
            final_savepath = save_path + '/' + filename[i] + '.png'
            im = Image.fromarray(final_mask)
            im.save(final_savepath)
   

if __name__ == '__main__':
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    main()